#!/usr/bin/python
# coding: utf-8

from xml.dom.minidom import parse

#
# Clase mensaje del Rafa para mantener el formato
#
class mensaje:
    def __init__(self):
        self.autor = ""
        self.fecha = ""
        self.titulo = ""
        self.id = ""
        self.en_respuesta_a = ""
        self.texto = ""
        self.permalink = ""
        

    def get_dom(self,doc):
        root = doc.createElement("mensaje")
        meta = doc.createElement("metadatos")
        texto = doc.createElement("texto")
        root.appendChild(meta)
        root.appendChild(texto)
        t = doc.createTextNode(self.texto)
        texto.appendChild(t)
        self._setattribute(meta,"autor")
        self._setattribute(meta,"fecha")
        self._setattribute(meta,"titulo")
        self._setattribute(meta,"id")
        self._setattribute(meta,"en_respuesta_a")
        self._setattribute(meta,"permalink")
        return root

    def _setattribute(self,son,varname):
        son.setAttribute(varname,self.__dict__[varname])

    def _loadvar(self,son,varname):
        if son.attributes.has_key(varname):
            self.__dict__[varname] = son.attributes[varname].value

    def load_from_dom(self,root):
        for son in root.childNodes:
            if son.nodeName == "metadatos":
                self._loadvar(son,"autor")
                self._loadvar(son,"fecha")
                self._loadvar(son,"titulo")
                self._loadvar(son,"id")
                self._loadvar(son,"en_respuesta_a")
                self._loadvar(son,"permalink")
            elif son.nodeName == "texto":
                if len(son.childNodes) != 0 and son.childNodes[0].nodeType == son.childNodes[0].TEXT_NODE:
                    self.texto = son.childNodes[0].data

import math
# Not defined in any module
def log2(val):
	return math.log10(val)/math.log10(2)

# Subclass of dict.
class freq(dict) :
	def N(self) :
		n = 0
		for k,v in self.iteritems() :
			n += v
		return n
	def prob_key(self, key) :
		return 	float(self[key])/self.N()
	def entropy(self) :
		h = 0
		for x in self.keys() :
			p_x = self.prob_key(x)
			h += p_x * log2(p_x)
		return -h
	def topK(self, k) :
		l = self.items()
		l.sort(lambda x,y : cmp(x[1],y[1]), reverse=True)
		l = l[:k]
		return l

class freqNGram(freq) :
	def __init__(self,n) :
		self.ngram = n
	def add_text(self, text) :
		for i in range(len(text)-n) :
			g = text[i:i+self.ngram+1]
			#print "DEBUG:", g
			if self.has_key(g) :
				self[g] += 1
			else :
				self[g] = 1

import sys

#
# Main
#

if len(sys.argv) > 1:
	assert(sys.argv[1])
	fd = open(sys.argv[1], "r")					# open file
	if len(sys.argv)>2 and sys.argv[2] : n = int(sys.argv[2])	# ngram
	else :	n = 0
	if len(sys.argv)>3 and sys.argv[3] : k = int(sys.argv[3])	# top-k
	else :	k = 20

	m = mensaje()
	f = freqNGram(n)
	doc = parse(fd)
	for root in doc.childNodes :
		if root.localName=="corpus" :
			for message in root.childNodes :
				#print "DEBUG: ", message.localName
				if message.localName=="mensaje" :
					m.load_from_dom(message)
					f.add_text(m.texto)
	fd.close()

	print "#",f.ngram,"-grams total", f.N()
	print "# Entropy", f.entropy()
	print "# Top", k
	toplist = f.topK(k)
	for k,v in toplist :
		print k,v
else:
	print "Usage:", sys.argv[0], "corpus.xml n-gram top-k"
	print
	print "Example:", sys.argv[0], "clarin.xml 2 50, gets frequency and entropy of top 50 2-grams"

